import logging
from datetime import datetime
from pathlib import Path
from timeit import default_timer as timer

import cv2
import numpy as np
from acllite.acllite_model import AclLiteModel
from acllite.acllite_resource import AclLiteResource
from utils import get_args, get_person, letterbox


class Detector:
    def __init__(self):
        self.infer_type = None
        self.acl_resource = AclLiteResource()
        self.acl_resource.init()
        self.model_path = get_args("model_path")
        self.out_classes = get_args("yolo_classes")["output"]
        self.yolo_model = AclLiteModel(self.model_path["yolo"])
        self.classify_model = None
        self.yolo_classes = None
        self.save_screenshot = False

        if get_args("screenshot")["save_screenshot"]:
            self.save_screenshot = True
            self.last_save_time = 0
            self.screenshot_save_interval = get_args("screenshot")[
                "screenshot_save_interval"
            ]
            self.img_base_path = Path(get_args("image_base_path")) / "screenshot"
            self.img_base_path = self.img_base_path / datetime.now().strftime(
                "%Y%m%d%H%M%S"
            )
            while self.img_base_path.exists():
                self.img_base_path = self.img_base_path.parent / (
                    self.img_base_path.name + "_1"
                )
            self.img_base_path.mkdir(parents=True)
            logging.info(f"Create image base path: {self.img_base_path}")

    def set_detector(self, infer_type):
        self.infer_type = infer_type
        self.yolo_classes = get_args("yolo_classes")[infer_type]
        if self.classify_model:
            self.classify_model.destroy()
        self.classify_model = AclLiteModel(self.model_path[f"classify_{infer_type}"])
        self.classify_classes = get_args("classify_classes")

    def detect(self, frame, neth=640, netw=640):
        height, width = frame.shape[:2]
        img = letterbox(frame, new_shape=(neth, netw))[0]
        img_info = np.array([neth, netw, height, width], dtype=np.float32)
        result = self.yolo_model.execute([img, img_info])  # net out, infer time
        batch_boxout, boxnum = result

        num_det = int(boxnum[0][0])
        boxout = (
            batch_boxout[0][: num_det * 6].reshape(6, -1).transpose().astype(np.float32)
        )
        # boxout: [[x1, y1, x2, y2, pres, cls], ...]
        other_boxout = [
            list(x[:4]) + [self.out_classes[int(x[5])]]
            for x in boxout
            if self.out_classes[int(x[5])] not in self.yolo_classes
            and self.infer_type == "touch"
        ]
        boxout = [x for x in boxout if self.out_classes[int(x[5])] in self.yolo_classes]
        objs = []
        if boxout is not None:
            for xyxy in boxout:
                xyxy = xyxy[:4]
                person = get_person(frame, xyxy)

                #person = person[..., ::-1]
                person = cv2.resize(person, (224, 224))
                classify_result = self.classify_model.execute([person])
                vals = classify_result[0].flatten()
                obj_class = vals.argsort()[-1]
                obj_class = self.reset_class(int(obj_class), self.infer_type)
                obj = list(xyxy)
                obj.append(self.classify_classes[obj_class])
                # print(self.classify_classes[obj_class])
                objs.append(obj)
        objs += other_boxout

        if self.save_screenshot:
            if timer() - self.last_save_time > self.screenshot_save_interval:
                self.last_save_time = timer()
                img_name = f"shot_{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.jpg"
                img_path = self.img_base_path / img_name
                frame = cv2.resize(frame, (640, 360))
                cv2.imwrite(str(img_path), frame, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
        return objs

    # FIXME put this into to code
    def reset_class(self, cls_num, infer_type):
        if infer_type == "sleep":
            if cls_num == 2 or cls_num == 3:
                return 1
            else:
                return 0
        if infer_type == "touch":
            # touch
            if cls_num == 0:
                return 3
            # untouch
            if cls_num == 1:
                return 4
            # unknow
            if cls_num == 2:
                return 2
            # if cls_num == 0:
            #     return 3
            # if cls_num == 1:
            #     return 4
            # if cls_num == 2:
            #     return 2
        return cls_num


class ArgsObj:
    def __init__(self, args):
        self.args = args

    def __getattr__(self, item):
        return self.args[item]
